package nameless.aistudy.common.util.ai.impl;

import nameless.aistudy.common.constant.ai.AIConstants;
import nameless.aistudy.common.util.ai.AIUtil;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;

@Component
public class AIUtilImpl implements AIUtil {

    ChatClient.Builder builder;

    public AIUtilImpl(ChatClient.Builder builder) {
        this.builder = builder;
    }

    @Override
    public ChatClient buildChatClient() {
        return buildChatClient(null, null);
    }

    @Override
    public ChatClient buildChatClient(Integer maxTokens, Double temperature) {
        if (temperature == null) temperature = AIConstants.DEFAULT_TEMPERATURE;
        if (maxTokens == null) maxTokens = AIConstants.DEFAULT_MAX_TOKENS;

        ChatOptions options = OpenAiChatOptions.builder()
                .maxTokens(maxTokens)
                .temperature(temperature)
                .build();

        return this.builder
                .defaultOptions(options)
                .build();
    }

    @Override
    public ChatClient.ChatClientRequestSpec buildChatClientRequest(String prompt) {
        return buildChatClient().prompt(prompt);
    }

    @Override
    public ChatClient.ChatClientRequestSpec buildChatClientRequest(String systemPrompt, String userPrompt) {
        return buildChatClient().prompt()
                .system(systemPrompt)
                .user(userPrompt);
    }

    @Override
    public SseEmitter buildSseEmitter(String prompt) {
        return buildSseEmitter(buildChatClientRequest(prompt));
    }

    @Override
    public SseEmitter buildSseEmitter(String systemPrompt, String userPrompt) {
        return buildSseEmitter(buildChatClientRequest(systemPrompt, userPrompt));
    }

    private SseEmitter buildSseEmitter(ChatClient.ChatClientRequestSpec chatClientRequest) {
        SseEmitter emitter = new SseEmitter();

        chatClientRequest
                .stream().content()
                .map(content -> SseEmitter.event().data(content))
                .subscribe(
                        event -> {
                            try {
                                emitter.send(event);  // 手动处理IOException
                            } catch (IOException e) {
                                // 客户端可能已断开连接
                                throw new RuntimeException("SSE Send failed", e);
                            }
                        },
                        emitter::completeWithError,
                        emitter::complete
                );

        return emitter;
    }

}
